"""
Copyright (c) 2022 Guangdong University of Technology
PhotLab is licensed under [Open Source License].
You can use this software according to the terms and conditions of the [Open Source License].
You may obtain a copy of [Open Source License] at: [https://open.source.license/]

THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.

See the [Open Source License] for more details.

Author: Meng Xiang, Junjiang Xiang, Sailan Yan
Created: 2023/8/19
Supported by: National Key Research and Development Program of China

"""

import phot


def QAM_of_transmitter_DSP(num_symbols, bits_per_symbol, up_sampling_factor, RRC_ROLL_OFF, total_baud):
    """ 首先产生发射端X/Y双偏振信号 """
    signal_bits = phot.gen_bits(num_symbols * bits_per_symbol)

    """# 生成双偏振符号序列 #"""
    signals = phot.modulation(signal_bits, bits_per_symbol)

    prev_symbols = signals

    """# 上采样 #"""
    signals = phot.up_sample(signals, up_sampling_factor)

    """# RRC #"""
    signals = phot.pulse_shaper(
        signals, up_sampling_factor, RRC_ROLL_OFF, total_baud)

    return signals, prev_symbols


def QAM_of_receiver_DSP(signals, prev_symbols, sampling_rate, up_sampling_factor, RRC_ROLL_OFF, total_baud, num_tap,
                        ref_power_cma, cma_convergence, step_size_cma, step_size_rde, N_Test_Angle, Block_Size):
    """ GSOP """
    signals = phot.iq_compensation(signals)

    """ 粗频偏估计和补偿 """
    signals = phot.freq_offset_compensation(signals, sampling_rate)

    signals = phot.pulse_shaper(
        signals, up_sampling_factor, RRC_ROLL_OFF, total_baud)

    # 同步 #
    signals, prev_symbols = phot.synchronization(
        signals, prev_symbols, up_sampling_factor)

    """ 自适应均衡/解偏振：采用CMA-RDE算法 """
    # num_tap = 25  # 均衡器抽头数目，此处均衡器内部是采用FIR滤波器，具体可查阅百度或者论文，
    # ref_power_cma = 2  # 设置CMA算法的模
    # cma_convergence = 30000  # CMA预均衡收敛的信号长度
    # step_size_cma = 1e-9  # CMA的更新步长，梯度下降法的步长
    # step_size_rde = 1e-9  # RDE的更新步长，梯度下降法的步长，%% CMA和RDE主要就是损失函数不同
    signals = phot.adaptive_equalize(
        signals,
        num_tap,
        cma_convergence,
        ref_power_cma,
        step_size_cma,
        step_size_rde,
        up_sampling_factor,
        bits_per_symbol,
        total_baud,
    )

    # 精确的频偏估计和补偿：采用FFT-FOE算法，
    signals = phot.freq_offset_compensation(signals, total_baud)

    """ 相位恢复：采用盲相位搜索算法进行相位估计和补偿 """
    # N_Test_Angle = 64       # BPS算法的测试角数目
    # Block_Size = 100        # BPS算法的块长设置
    signals = phot.bps_restore(
        signals, N_Test_Angle, Block_Size, bits_per_symbol)
    phot.constellation_diagram(signals)

    """ 同步 """
    signals, prev_symbols = phot.synchronization(signals, prev_symbols, 1)

    """ 相位校正，可供选择 """
    signals, prev_symbols = phot.Correct_phase(prev_symbols, signals)

    """ 解码 """
    signals = phot.deModulation(signals, bits_per_symbol)
    prev_symbols = phot.deModulation(prev_symbols, bits_per_symbol)

    return signals, prev_symbols


if __name__ == '__main__':

    phot.config(plot=True, backend="torch")

    """ 系统参数 """
    num_symbols = 2**16                            # 符号数目
    bits_per_symbol = 3                            # 2 for QPSK; 4 for 16QAM; 6 for 64QAM  设置调制格式
    total_baud = 40e9                              # 信号波特率，符号率
    up_sampling_factor = 4                         # 上采样倍数
    sampling_rate = up_sampling_factor * total_baud  # 信号采样率
    Reference_frequency = 193.1e12
    """ 发射端 """
    """ 首先产生发射端X/Y双偏振信号 """
    signal_bits = phot.gen_bits(num_symbols, bits_per_symbol)

    """ 生成双偏振符号序列 """
    signals = phot.modulation(signal_bits, bits_per_symbol)

    prev_symbols = signals

    """ 上采样 """
    signals = phot.up_sample(signals, up_sampling_factor)

    """ RRC """
    RRC_ROLL_OFF = 0.01
    signals = phot.pulse_shaper(signals, up_sampling_factor, RRC_ROLL_OFF, total_baud)

    """ 链路 """
    """ OSNR """
    osnr = 32  # 设置系统OSNR，也就是光信号功率与噪声功率的比值，此处单位为dB
    signals = phot.gaussian_noise(signals, osnr, sampling_rate)

    """ 接收端 """
    """ GSOP """
    signals = phot.iq_compensation(signals)

    """ 粗频偏估计和补偿 """
    signals = phot.freq_offset_compensation(signals, sampling_rate)

    signals = phot.pulse_shaper(signals, up_sampling_factor, RRC_ROLL_OFF, total_baud)

    """ 同步 """
    signals, prev_symbols = phot.synchronization(signals, prev_symbols, up_sampling_factor)

    """ 自适应均衡/解偏振: 采用CMA-RDE算法 """
    num_tap = 25              # 均衡器抽头数目，此处均衡器内部是采用FIR滤波器，具体可查阅百度或者论文，
    ref_power_cma = 2         # 设置CMA算法的模
    cma_convergence = 33000   # CMA预均衡收敛的信号长度
    step_size_cma = 2e-9      # CMA的更新步长，梯度下降法的步长
    step_size_rde = 2e-9      # RDE的更新步长，梯度下降法的步长，%% CMA和RDE主要就是损失函数不同
    signals = phot.adaptive_equalizer(signals, num_tap, cma_convergence, ref_power_cma,step_size_cma, step_size_rde, up_sampling_factor, bits_per_symbol, total_baud)

    """精确的频偏估计和补偿: 采用FFT-FOE算法"""
    signals = phot.freq_offset_compensation(signals, sampling_rate)

    """ 相位恢复：采用盲相位搜索算法进行相位估计和补偿 """
    num_test_angle = 64       # BPS算法的测试角数目
    block_size = 120        # BPS算法的块长设置
    signals = phot.bps_restore(signals, num_test_angle, block_size, bits_per_symbol)

    """分析器画星座图"""
    phot.constellation_diagram(signals, isdata = False)


